Skip to content

Parallelize and right-size the sequence-properties compute step#22

Draft
PaulNewling wants to merge 6 commits into
mainfrom
chore/deterministic-threads
Draft

Parallelize and right-size the sequence-properties compute step#22
PaulNewling wants to merge 6 commits into
mainfrom
chore/deterministic-threads

Conversation

@PaulNewling

Copy link
Copy Markdown
Collaborator

Summary

The compute step ran single-threaded under a flat mem("12GiB") / cpu(1), leaving cores idle and risking OOM on large datasets. This PR parallelizes it, sizes its memory and CPU to the input, and speeds up the compute engine — all while keeping output byte-identical (the quantized-equal / CID-stability contract).

At 1M rows: peptide compute ~2.3× faster; antibody wall time ~10.0s → ~4.6s (it now scales with cores, 1.1× → 1.5× on 4).

What changed

Deterministic multithreading. main.py pins the BLAS/OpenMP intra-op threads to 1, so the one order-sensitive step (the counts @ weights reduction) never splits across threads. Polars then parallelizes only order-stable work (CSV read, chain reconstruction, the aa_fraction reshape, the unique-key sort, CSV write). The emitted bytes are independent of core count.

Input-scaled resources (replacing the flat mem/cpu):

  • memFormula(clamp(max(lineCount·perRow, size·64), 2GiB, 64GiB)) — a row-scaled term for per-clone structures (count matrix, output columns, aa_fraction) plus a residue-scaled size term for the O(residues) cleaning transients. The size term keeps long-sequence/low-row inputs (amplicon variants) off the 2 GiB floor. Constants are calibrated from measured peak RSS.
  • cpuFormula scales both modes up to 4 cores. POLARS_MAX_THREADS is set to the cores the backend actually grants, via the {system.cpu} command expression, so a sub-cap allocation never oversubscribes its quota.
  • Both carry fallbacks for backends without formula support (fallback: "12GiB", static cpu/env).

Compute-engine optimizations (all output-preserving):

  • _charge_raw computes 10**ph once and reuses it via scalar factors (~2.4× on the charge/pI path).
  • Peptide mode cleans its column once, sharing the intermediate between the count substrate and the instability index.
  • The count and instability scatters use np.bincount instead of the slower np.add.at (~3×, bit-identical).
  • The median-CDR3-length stat uses a vectorized polars count_matches instead of one Python call per row.
  • Peptide aa_fraction uses a polars unpivot instead of a 20×N Python-list explosion, cutting peak memory.
  • The pI bisection runs row-parallel across threads (contiguous, data-independent blocks; bit-identical regardless of worker count).

Determinism

Output stays within the quantized-equal contract, verified byte-for-byte:

  • new vs the pre-change implementation — peptide and antibody, including edge cases (empty / stop-codon / lowercase / non-standard / missing-region);
  • across 1 vs 4 threads at 1M rows (guards the parallel bisection).

The charge refactor drifts ~1e-15, below the 3-dp quantization. The one theoretical exception is a bisection midpoint whose net charge lands within ~1e-15 of zero (astronomically rare, never observed).

Testing

  • 262 software tests pass (unit + integration), including an extended determinism test that runs the CLI at 1 vs 4 threads and a ≥50k-row case that exercises the parallel bisection.
  • ruff check + ruff format clean.
  • pnpm build:dev green (7/7 tasks; the workflow Tengo compiles).
  • Benchmarked peak RSS and thread scaling to calibrate the formula constants.

Not yet verified

The resource formulas and {system.cpu} resolve at execution time, so the runtime path (the lineCount pre-exec → computed quota, and the command-expression substitution) needs a live-backend run to exercise end-to-end. The Tengo compiles and the mechanism is confirmed against the backend source, but it has not yet run against a live server.

Scale the compute step with the dataset while keeping output byte-identical (the quantized-equal / CID-stability contract).

Deterministic threading: the workflow requests cores and passes them to polars via POLARS_MAX_THREADS; main.py pins BLAS/OpenMP intra-op threads to 1 so the order-sensitive 'counts @ weights' reduction never splits across threads. Emitted bytes are independent of core count.

Input-scaled resources: replace the flat mem/cpu with exec resource formulas — a lineCount-based memFormula for both modes (fallback 12GiB), and a cpuFormula up to 4 cores for peptide/amplicon (antibody stays single-core; it is ~1.1x on threads). Constants calibrated from measured peak RSS (~3 KiB/row peptide, ~4.3 KiB/row antibody) and thread scaling.

Compute engine: _charge_raw computes 10**ph once and reuses it via scalar factors (~2.4x on the charge/pI path); peptide cleans its column once and shares the intermediate between counts and instability; count/instability scatters use np.bincount instead of np.add.at (~3x). aa_fraction long-format uses a polars unpivot instead of a 20xN Python-list explosion. Peptide compute ~2.3x, antibody wall ~25% at 1M rows.

Verified: 261 software tests pass; output byte-identical vs the pre-change implementation for peptide and antibody (incl. empty/stop-codon/lowercase/non-standard/missing-region edge cases) and across 1 vs 4 polars threads.

Bumps @platforma-sdk/workflow-tengo to 6.6.5 (catalog) for the resource-formula API.
…stat

Second round of compute-engine speedups on top of the previous commit; output stays byte-identical (quantized-equal / CID-stability contract).

Vectorize the per-chain median-CDR3-length stat (_median_cdr3_length_by_chain) with a polars count_matches over standard-AA chars instead of one effective_length Python call per row — the largest antibody-mode stats cost. Median formula unchanged, so stats.json is byte-identical.

Row-parallelize the pI bisection (isoelectric_point / fv_isoelectric_point): contiguous, data-independent row blocks bisected on separate threads (numpy elementwise releases the GIL), concatenated in row order. Each row's bisection depends only on its own counts, so the result is bit-identical regardless of worker count.

With both, antibody now scales with cores (~1.13x -> ~1.53x on 4 cores; 10.0s -> 4.6s at 1M rows), so the workflow scales antibody CPU up to 4 as well (it was pinned to cpu(1) when the mode did not parallelize). Both modes now share the same cpuFormula.

Add a determinism test above vectorized._PARALLEL_MIN_ROWS (antibody, exercising both isoelectric_point and fv_isoelectric_point) asserting byte-identical output at 1 vs 4 threads; the smaller cases stay under the parallel threshold and cannot reach it.

Verified: 262 software tests pass; output byte-identical vs the pre-change implementation (peptide + antibody at 1M) and across 1 vs 4 threads at scale.
…equence OOM

The lineCount-only memFormula under-provisioned low-row/long-sequence inputs: peak RSS is dominated by O(total-residues) cleaning transients (the DIWV weight gather, residue-row ids, byte buffers), which lineCount does not see. An amplicon-like 200k-row x ~600aa input peaks at 4.51 GiB but the old formula sized it to 200k*5KiB -> the 2 GiB floor, an OOM the pre-change flat 12 GiB avoided.

Add a residue-scaled term size("input.tsv") * 64 (measured ~40 B peak RSS per input byte; 64x for headroom) and take max() with the row-scaled lineCount term, so each dominates in its regime: short sequences use the row term (no change), long sequences use the size term. Verified: amplicon case now sizes to ~7.25 GiB >= 4.51 GiB peak; normal peptide/antibody sizing is unchanged (row term still dominates).
Review follow-ups #3 and #4, both byte-identical (no numeric or output change).

#3: _charge_raw takes an optional precomputed ten_ph; fv_isoelectric_point computes 10**ph once and shares it across both chains instead of recomputing the array-wide transcendental twice per Fv bisection step. Same value -> bit-identical.

#4: build the median-CDR3 count_matches character class from STANDARD_AAS (_STANDARD_AA_CLASS) instead of a hardcoded 40-char literal, so it cannot drift from properties.effective_length / the count substrate. Same class -> bit-identical.

Also softened the _charge_raw docstring's absolute 'bytes/CID unchanged' to 'within the quantized-equal contract' (review #5): the FP drift is byte-identical on the corpus but a knife-edge bisection midpoint within ~1e-15 of zero is a rare theoretical exception.

Verified: 262 tests pass, ruff clean, output byte-identical vs prior HEAD on antibody 1M (exercises both the Fv pI path and the median-CDR3 stat).
…t the static cap

Resolves review #2 (thread oversubscription). The workflow pinned POLARS_MAX_THREADS to the static cap (computeCpu=4) while cpuFormula grants 1..4 cores dynamically, and vectorized._n_workers() reads POLARS_MAX_THREADS to size the pI-bisection pool. So a sub-cap input (>=50k but <250k rows -> 1 core granted) spawned 4 polars + 4 bisection threads on a 1-core quota.

Set it instead with envExpr("POLARS_MAX_THREADS", "{system.cpu}"): the backend's command-expression engine substitutes the resolved quota grant (ComputeLimitsInfo.CPU = the cpuFormula result) before the process starts, so main.py reads the true granted core count at import. Gated on feats.commandExpressions with a static env(string(computeCpu)) fallback for older backends. This also removes the leaky reuse of the static cap and keeps the thread count in lockstep with the CPU reservation.

Verified: workflow tengo compiles (feats/envExpr), 262 tests pass. The runtime {system.cpu} substitution resolves at execution time, so it needs the live-backend integration test to exercise end-to-end.
Changeset: split the memory/CPU/threads wall into one-idea paragraphs (memory, CPU, determinism), convert the compute-optimization run-in list to scannable bullets, cut wordiness, and drop repeated byte-identity claims.

Comments: make the aa_fraction build comment active voice ('is unpivoted'/'are made contiguous' -> active), and fix a stale _n_workers docstring (POLARS_MAX_THREADS is the granted cores now, not the static cpu() request). No code changes; ruff + tests green.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces parallelization and resource scaling to the compute step, optimizing performance while maintaining strict output determinism. Key optimizations include cleaning columns once, replacing slow np.add.at scatters with np.bincount, vectorizing CDR3 length calculations, and restructuring the peptide amino acid fraction calculation using a Polars unpivot. To guarantee byte-identical outputs, BLAS/OpenMP intra-op threads are pinned to 1. Feedback on these changes highlights opportunities to improve performance and type safety: removing a redundant np.ascontiguousarray call in the peptide path, correcting incomplete type annotations for ph and ten_ph in _charge_raw, and adding missing type annotations to the new _bisect_rows_parallel helper.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread software/src/pipeline.py
Comment on lines +237 to +240
wide = pl.DataFrame(
{"entity_key": key_series}
| {aa: _na_to_null(np.ascontiguousarray(fractions[:, i])) for i, aa in enumerate(STANDARD_AAS)}
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The use of np.ascontiguousarray here is redundant. Since _na_to_null iterates over the elements of the column slice to convert NaN values to None (resulting in a Python list or object array), the memory contiguity of the input numpy array does not provide any performance benefit. Calling np.ascontiguousarray creates an unnecessary intermediate copy of the float array slice for each of the 20 columns, which wastes CPU cycles and memory allocations.

Suggested change
wide = pl.DataFrame(
{"entity_key": key_series}
| {aa: _na_to_null(np.ascontiguousarray(fractions[:, i])) for i, aa in enumerate(STANDARD_AAS)}
)
wide = pl.DataFrame(
{"entity_key": key_series}
| {aa: _na_to_null(fractions[:, i]) for i, aa in enumerate(STANDARD_AAS)}
)

Comment on lines +349 to +351
def _charge_raw(
sub: Substrate, ph: float, pka_set: PKaSet, include_cys: bool, ten_ph: np.ndarray | None = None
) -> np.ndarray:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The type annotations for ph and ten_ph are incomplete and can cause type-checker (e.g., MyPy) errors:

  1. ph is annotated as float, but during the vectorized bisection, it is passed as a numpy array (np.ndarray).
  2. ten_ph is annotated as np.ndarray | None, but if ph is a scalar float and ten_ph is not provided, ten_ph is computed as 10.0**ph (which is a float). Assigning a float to a variable annotated as np.ndarray | None will trigger a type-checker warning.

Updating the annotations to support both scalar and array types resolves these issues.

def _charge_raw(
    sub: Substrate,
    ph: float | np.ndarray,
    pka_set: PKaSet,
    include_cys: bool,
    ten_ph: float | np.ndarray | None = None,
) -> np.ndarray:

return 1


def _bisect_rows_parallel(make_charge_fn, valid, lo=0.0, hi=14.0, tol=1e-3):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The helper function _bisect_rows_parallel is missing type annotations, whereas the rest of the file consistently uses type annotations. Adding type annotations improves code readability, maintainability, and type safety.

Note: This requires importing Callable from typing (i.e., from typing import Callable).

def _bisect_rows_parallel(
    make_charge_fn: Callable[[int, int], Callable[[float | np.ndarray], np.ndarray]],
    valid: np.ndarray,
    lo: float = 0.0,
    hi: float = 14.0,
    tol: float = 1e-3,
) -> np.ndarray:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant